package telex;
import com.aliasi.annotate.corpora.HtmlToText;
import com.aliasi.tokenizer.IndoEuropeanTokenCategorizer;
import com.aliasi.tokenizer.IndoEuropeanTokenizerFactory;
import com.aliasi.tokenizer.Tokenizer;
import com.sleepycat.bind.EntryBinding;
import com.sleepycat.bind.serial.SerialBinding;
import com.sleepycat.bind.tuple.IntegerBinding;
import com.sleepycat.bind.tuple.StringBinding;
import com.sleepycat.bind.serial.StoredClassCatalog;
import com.sleepycat.collections.StoredEntrySet;
import com.sleepycat.collections.StoredIterator;
import com.sleepycat.collections.StoredKeySet;
import com.sleepycat.collections.StoredMap;
import com.sleepycat.collections.StoredSortedMap;
import com.sleepycat.je.Database;
import com.sleepycat.je.DatabaseConfig;
import com.sleepycat.je.DatabaseException;
import com.sleepycat.je.Environment;
import com.sleepycat.je.EnvironmentConfig;
import java.io.*;
import java.util.*;
import java.util.Map.Entry;
import java.util.zip.*;

public class NGramDBFromText {
    private static StoredSortedMap<String, Integer> ngramsDB;

    public NGramDBFromText(String dirNgramPath) throws DatabaseException {
        File dirNgram = new File(dirNgramPath);
        EnvironmentConfig envConfig = new EnvironmentConfig();
        envConfig.setTransactional(false);
        envConfig.setAllowCreate(true);
        Environment env = new Environment(dirNgram, envConfig);
        DatabaseConfig dbConfig = new DatabaseConfig();
        dbConfig.setTransactional(false);
        dbConfig.setAllowCreate(true);
        dbConfig.setSortedDuplicates(false);
        dbConfig.setDeferredWrite(true);
        final Database db = env.openDatabase(null, "ngramsdatabase", dbConfig);
        final StoredClassCatalog catalog = new StoredClassCatalog(env.openDatabase(null, "ngramsdatabase-cat", dbConfig));
        EntryBinding<String> keyBinding = new SerialBinding(catalog, ((String) (new String())).getClass());
        EntryBinding<Integer> valueBinding = new SerialBinding(catalog, ((Integer) (new Integer(0))).getClass());
        this.ngramsDB = new StoredSortedMap<String, Integer>(db, keyBinding, valueBinding, true);
        Runtime.getRuntime().addShutdownHook( new Thread( ) {
           public void run() {
                       try {
                           db.sync();
                           db.close();
                           catalog.close();
                           db.getEnvironment().sync();
                           db.getEnvironment().cleanLog();
                           db.getEnvironment().close();
                       } catch ( DatabaseException ex ) {
                               ex.printStackTrace();
                       }
           }
       });
    }

    private static void addGram(String gram) {
        if(!ngramsDB.containsKey(gram))
            ngramsDB.put(gram, 1);
        else
            ngramsDB.put(gram, ngramsDB.get(gram) + 1);
    }

    private static long extractGramsFromFile(File inFile, int gramSize) throws
            Exception {
            System.out.println("Ficheiro: " + inFile);
        BufferedReader brIn = new BufferedReader(new FileReader(inFile));
        String sLine;
        int total = 0;
        while ((sLine = brIn.readLine()) != null) {
            if(!sLine.equals("")){
                Tokenizer tok = IndoEuropeanTokenizerFactory.INSTANCE.tokenizer(sLine.toCharArray(), 0, sLine.length());
                String tokens[] = tok.tokenize();
                String words[] = new String[tokens.length + 2];
                words[0] = "<S>";
                System.arraycopy(tokens, 0, words, 1, tokens.length);
                words[words.length - 1] = "</S>";
                total += words.length;
                for (int i = 0; i < words.length; i++) {
                    String gram = words[i];
                    addGram(gram);
                    //ciclo para criar o n-gramas superior a uni-gramas
                    for (int j = 1; (j < gramSize && i + j < words.length); j++) {
                        gram += " " + words[i + j];
                        addGram(gram);
                    }
                }
            }
        }
        brIn.close();
        return total;
    }

    private static int getGramCount(String s) {
        int c = 1;
        int last = -1;
        while((last = s.indexOf(" ", last+1))>0) c++;
        return c;
    }

    public static void createNGramDB(String inDirectory, String outDirectory, int maxGramSize)
    throws Exception {
        File fIn = new File(inDirectory);
        File [] fInFiles = fIn.listFiles();
        String [] firstGrams = new String[maxGramSize];
        String s;
        long total = 0;
        for(int i=0; i<fInFiles.length; i++) {
            //System.out.println("A avaliar o ficheiro " + fInFiles[i].getName());
            total += extractGramsFromFile(fInFiles[i], maxGramSize);
        }
        System.out.println("passei os ngrams para a bd");
        ListIterator<String> it = (ListIterator<String>)(ngramsDB.keySet().iterator());
       
        File outDir = new File(outDirectory);
        GZIPOutputStream [] outStreams = new GZIPOutputStream[maxGramSize]; //1-grams go in outFiles[0]

        for (int i = 0; i < maxGramSize; i++) {
            String dirName = outDirectory + File.separator + (i + 1) + "gms";
            (new File(dirName)).mkdir();
            if (i == 0)
                outStreams[i] = new GZIPOutputStream(new FileOutputStream(dirName + File.separator + "vocab.gz"));
            else
                outStreams[i] = new GZIPOutputStream(new FileOutputStream(dirName + File.separator + (i+1) + "gm-0000.gz"));
        }
        System.out.println("escrevi os unigramas");
        outDir.mkdir();
        while(it.hasNext()) {
            String gram = it.next();
            int count = ngramsDB.get(gram);
            int numGrams = getGramCount(gram);
            if(firstGrams[numGrams-1] == null)
                firstGrams[numGrams-1] = gram;
            outStreams[numGrams - 1].write((gram + "\t" + count + "\n").getBytes());
        }

        for(int i=0; i< maxGramSize; i++) {
            outStreams[i].close();
            if(i==0) continue;
            BufferedWriter bwOut = new BufferedWriter(new FileWriter(outDirectory +
            		File.separator + (i + 1) + "gms" + File.separator + (i + 1) + "gm.idx"));
            new PrintWriter(bwOut).println((i+1) + "gm-0000.gz\t" + firstGrams[i]);
            bwOut.close();
        }
        BufferedWriter bwOut = new BufferedWriter(new FileWriter(outDirectory + File.separator +
                "1gms" + File.separator + "total"));
        bwOut.write(Long.toString(total));
        bwOut.close();
    }

    /*
     * Código antigo quando para criar os ngramas estoirava a memoria
     */
 /*
    private static void addGram(TreeMap<String, Integer> hm, String gram) {
        if(!hm.containsKey(gram))
            hm.put(gram, 1);
        else
            hm.put(gram, hm.get(gram)+1);
    }

    private static long extractGramsFromFile(File inFile, TreeMap<String,
                                             Integer> hm, int gramSize) throws
            Exception {

        BufferedReader brIn = new BufferedReader(new FileReader(inFile));
        String sLine;
        int total = 0;
        while ((sLine = brIn.readLine()) != null) {
            if(!sLine.equals("")){

                //String texto = "<S> " + sLine + " </S>";
                Tokenizer tok = IndoEuropeanTokenizerFactory.INSTANCE.tokenizer(sLine.toCharArray(), 0, sLine.length());
                String tokens[] = tok.tokenize();
                String words[] = new String[tokens.length + 1];
                words[0] = "<S>";
                System.arraycopy(tokens, 0, words, 1, tokens.length);
                //String[] words = ("<S> " + sLine + " </S>").split("[,() ?!]+");
                total += words.length;
                for (int i = 0; i < words.length; i++) {
                    String gram = words[i];
                    addGram(hm, gram);
                    //ciclo para criar o n-gramas superior a uni-gramas
                    for (int j = 1; (j < gramSize && i + j < words.length); j++) {
                        gram += " " + words[i + j];
                        addGram(hm, gram);
                    }
                }
            }
        }
        brIn.close();
        return total;
    }

    private static int getGramCount(String s) {
        int c = 1;
        int last = -1;
        while((last = s.indexOf(" ", last+1))>0) c++;
        return c;
    }

    public static void createNGramDB(String inDirectory, String outDirectory, int maxGramSize)
    throws Exception {
        File fIn = new File(inDirectory);
        File [] fInFiles = fIn.listFiles();
        TreeMap<String, Integer> grams = new TreeMap<String, Integer>();
        String [] firstGrams = new String[maxGramSize];
        long total = 0;
        for(int i=0; i<fInFiles.length; i++) {
            System.out.println("A avaliar o ficheiro " + fInFiles[i].getName());
            total += extractGramsFromFile(fInFiles[i], grams, maxGramSize);
        }
        Iterator<String> it = grams.keySet().iterator();
        File outDir = new File(outDirectory);
        GZIPOutputStream [] outStreams = new GZIPOutputStream[maxGramSize]; //1-grams go in outFiles[0]

        for (int i = 0; i < maxGramSize; i++) {
            String dirName = outDirectory + File.separator + (i + 1) + "gms";
            (new File(dirName)).mkdir();
            if (i == 0)
                outStreams[i] = new GZIPOutputStream(new FileOutputStream(dirName + File.separator + "vocab.gz"));
            else
                outStreams[i] = new GZIPOutputStream(new FileOutputStream(dirName + File.separator + (i+1) + "gm-0000.gz"));
        }
        outDir.mkdir();
        while(it.hasNext()) {
            String gram = it.next();
            int count = grams.get(gram);
            int numGrams = getGramCount(gram);
            if(firstGrams[numGrams-1] == null)
                firstGrams[numGrams-1] = gram;
            outStreams[numGrams - 1].write((gram + "\t" + count + "\n").getBytes());
        }

        for(int i=0; i< maxGramSize; i++) {
            outStreams[i].close();
            if(i==0) continue;
            BufferedWriter bwOut = new BufferedWriter(new FileWriter(outDirectory +
            		File.separator + (i + 1) + "gms" + File.separator + (i + 1) + "gm.idx"));
            new PrintWriter(bwOut).println((i+1) + "gm-0000.gz\t" + firstGrams[i]);
            bwOut.close();
        }
        BufferedWriter bwOut = new BufferedWriter(new FileWriter(outDirectory + File.separator +
                "1gms" + File.separator + "total"));
        bwOut.write(Long.toString(total));
        bwOut.close();
    }

    /*
     * Código novo para ver senão estoira a memoria
     */
/*
    private static void addGram(TreeMap<String, Integer> hm, String gram) {
        if(!hm.containsKey(gram))
            hm.put(gram, 1);
        else
            hm.put(gram, hm.get(gram)+1);
    }

    private static long extractGramsFromFile(File inFile, TreeMap<String,
                                             Integer> hm, int gramSize) throws
            Exception {

        BufferedReader brIn = new BufferedReader(new FileReader(inFile));
        String sLine;
        int total = 0;
        while ((sLine = brIn.readLine()) != null) {
            if(!sLine.equals("")){
                Tokenizer tok = IndoEuropeanTokenizerFactory.INSTANCE.tokenizer(sLine.toCharArray(), 0, sLine.length());
                String tokens[] = tok.tokenize();
                String words[] = new String[tokens.length + 2];
                words[0] = "<S>";
                System.arraycopy(tokens, 0, words, 1, tokens.length);
                words[words.length - 1] = "</S>";
                //String[] words = ("<S> " + sLine + " </S>").split("[,() ?!]+");
                total += words.length;
                for (int i = 0; i < words.length; i++) {
                    String gram = words[i];
                    if(gramSize == 0){
                        addGram(hm, gram);
                    }else if(i + gramSize < words.length){
                        for(int j = i+1; j-i<=gramSize; j++){
                           gram += " " + words[j];                           
                        }
                        addGram(hm, gram);   
                    }
                }
            }
        }
        brIn.close();
        return total;
    }

    private static int getGramCount(String s) {
        int c = 1;
        int last = -1;
        while((last = s.indexOf(" ", last+1))>0) c++;
        return c;
    }

    public static void createNGramDB(String inDirectory, String outDirectory, int maxGramSize)
    throws Exception {
        File fIn = new File(inDirectory);
        File [] fInFiles = fIn.listFiles();
        TreeMap<String, Integer> grams = new TreeMap<String, Integer>();
        String [] firstGrams = new String[maxGramSize];
        long total = 0;
        File outDir = new File(outDirectory);
        GZIPOutputStream [] outStreams = new GZIPOutputStream[maxGramSize];
        for(int j=0; j < maxGramSize; j++){
            System.out.println("Outro ngrama");
            for(int i=0; i<fInFiles.length; i++) {
                total += extractGramsFromFile(fInFiles[i], grams, j);
            }
            Iterator<String> it = grams.keySet().iterator();

            String dirName = outDirectory + File.separator + (j + 1) + "gms";
            (new File(dirName)).mkdir();
            if (j == 0)
            	outStreams[j] = new GZIPOutputStream(new FileOutputStream(dirName + File.separator + "vocab.gz"));
            else
            	outStreams[j] = new GZIPOutputStream(new FileOutputStream(dirName + File.separator + (j+1) + "gm-0000.gz"));

            outDir.mkdir();
            while(it.hasNext()) {
                String gram = it.next();
		int count = grams.get(gram);
		int numGrams = getGramCount(gram);
                if(firstGrams[numGrams-1] == null)
                    firstGrams[numGrams-1] = gram;
		outStreams[numGrams - 1].write((gram + "\t" + count + "\n").getBytes());
            }

            outStreams[j].close();
            grams.clear();
            if(j==0) continue;
            BufferedWriter bwOut = new BufferedWriter(new FileWriter(outDirectory +
		File.separator + (j + 1) + "gms" + File.separator + (j + 1) + "gm.idx"));
            new PrintWriter(bwOut).println((j+1) + "gm-0000.gz\t" + firstGrams[j]);
            bwOut.close();

        }

        BufferedWriter bwOut = new BufferedWriter(new FileWriter(outDirectory + File.separator +
            	"1gms" + File.separator + "total"));
        bwOut.write(Long.toString(total));
        bwOut.close();

    }
*/


    public static void main(String[] args) throws Exception {
        String inDir = null;
        String outDir = null;
        int gramSize = 0;
        try {
            Scanner scan = new Scanner(System.in);
            inDir = scan.nextLine();
            System.out.println("Directorio in: " + inDir);
            outDir = scan.nextLine();
            System.out.println("Directorio out: " + outDir);
            gramSize = Integer.parseInt(scan.nextLine());
            System.out.println("Ngram: " + gramSize);
            scan.close();
            //inDir = args[0];
            //outDir = args[1];
            //gramSize = Integer.parseInt(args[2]);
        }
        catch(Exception e) {
            System.out.println("Usage: NGramDBFromText <inDirectory> <outDirectory> <maxGramSize>");
            return;
        }
        if(gramSize == 0)
            return;
        System.out.println("vou criar");
        NGramDBFromText ndb = new NGramDBFromText("c:/users/diogo/desktop/ngrams_db");
        ndb.createNGramDB(inDir, outDir, gramSize);
    }
}
